#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
数据集管理模块
"""

import os
import json
import shutil
import datetime
from pathlib import Path

class DatasetManager:
    """数据集管理类"""
    
    def __init__(self, base_dir, logger):
        """
        初始化数据集管理器
        
        Args:
            base_dir (str): 数据集基础目录
            logger (logging.Logger): 日志器
        """
        self.base_dir = Path(base_dir)
        self.logger = logger
        self.current_dataset = None
        
        # 确保目录存在
        if not self.base_dir.exists():
            self.base_dir.mkdir(parents=True, exist_ok=True)
    
    def create_dataset(self, name, algorithm_type):
        """
        创建新数据集
        
        Args:
            name (str): 数据集名称
            algorithm_type (str): 算法类型（如 'classification', 'detection'）
            
        Returns:
            dict: 包含数据集信息的字典
        """
        # 生成唯一数据集ID和时间戳
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        dataset_id = f"{name}_{timestamp}"
        
        # 创建数据集目录
        dataset_path = self.base_dir / dataset_id
        if dataset_path.exists():
            self.logger.warning(f"数据集路径已存在：{dataset_path}")
            return None
        
        # 创建子目录
        images_dir = dataset_path / "images"
        annotations_dir = dataset_path / "annotations"
        
        dataset_path.mkdir(parents=True, exist_ok=True)
        images_dir.mkdir(parents=True, exist_ok=True)
        annotations_dir.mkdir(parents=True, exist_ok=True)
        
        # 创建元数据
        metadata = {
            "id": dataset_id,
            "name": name,
            "created_at": timestamp,
            "algorithm_type": algorithm_type,
            "image_count": 0,
            "annotation_count": 0,
            "labels": []
        }
        
        # 保存元数据
        with open(dataset_path / "dataset_meta.json", "w", encoding="utf-8") as f:
            json.dump(metadata, f, ensure_ascii=False, indent=2)
        
        self.current_dataset = metadata
        self.current_dataset["path"] = str(dataset_path)
        
        self.logger.info(f"已创建数据集：{name} ({dataset_id})")
        return self.current_dataset
    
    def load_dataset(self, dataset_path):
        """
        加载现有数据集
        
        Args:
            dataset_path (str): 数据集路径
            
        Returns:
            dict: 包含数据集信息的字典，加载失败则返回None
        """
        dataset_path = Path(dataset_path)
        metadata_file = dataset_path / "dataset_meta.json"
        
        if not metadata_file.exists():
            self.logger.error(f"数据集元数据不存在：{metadata_file}")
            return None
        
        try:
            with open(metadata_file, "r", encoding="utf-8") as f:
                metadata = json.load(f)
            
            # 添加路径信息
            metadata["path"] = str(dataset_path)
            self.current_dataset = metadata
            
            self.logger.info(f"已加载数据集：{metadata['name']} ({metadata['id']})")
            return self.current_dataset
            
        except Exception as e:
            self.logger.error(f"加载数据集失败：{e}")
            return None
    
    def get_datasets_list(self):
        """
        获取所有数据集的列表
        
        Returns:
            list: 数据集信息的列表
        """
        datasets = []
        
        for item in self.base_dir.iterdir():
            if item.is_dir():
                metadata_file = item / "dataset_meta.json"
                if metadata_file.exists():
                    try:
                        with open(metadata_file, "r", encoding="utf-8") as f:
                            metadata = json.load(f)
                        
                        # 添加路径信息
                        metadata["path"] = str(item)
                        datasets.append(metadata)
                    except Exception as e:
                        self.logger.warning(f"读取数据集元数据失败：{e}")
        
        # 按创建时间排序
        datasets.sort(key=lambda x: x["created_at"], reverse=True)
        return datasets
    
    def add_images(self, image_paths, copy=True):
        """
        向当前数据集添加图像
        
        Args:
            image_paths (list): 图像文件路径列表
            copy (bool): 是否复制文件（True）或移动文件（False）
            
        Returns:
            int: 添加的图像数量
        """
        if not self.current_dataset:
            self.logger.error("没有打开的数据集")
            return 0
        
        dataset_path = Path(self.current_dataset["path"])
        images_dir = dataset_path / "images"
        count = 0
        
        for img_path in image_paths:
            img_path = Path(img_path)
            if not img_path.exists() or not img_path.is_file():
                self.logger.warning(f"图像不存在或不是文件：{img_path}")
                continue
            
            # 检查文件类型
            if img_path.suffix.lower() not in ['.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff']:
                self.logger.warning(f"不支持的图像格式：{img_path}")
                continue
            
            # 生成目标路径
            target_path = images_dir / img_path.name
            
            # 如果文件已存在，添加数字后缀
            if target_path.exists():
                base_name = target_path.stem
                extension = target_path.suffix
                i = 1
                while target_path.exists():
                    target_path = images_dir / f"{base_name}_{i}{extension}"
                    i += 1
            
            # 复制或移动文件
            try:
                if copy:
                    shutil.copy2(img_path, target_path)
                else:
                    shutil.move(img_path, target_path)
                count += 1
            except Exception as e:
                self.logger.error(f"添加图像失败：{e}")
        
        # 更新元数据
        if count > 0:
            self.current_dataset["image_count"] += count
            self._save_metadata()
            self.logger.info(f"已添加{count}张图像到数据集")
        
        return count
    
    def add_label(self, name, color="#FF0000"):
        """
        添加标签
        
        Args:
            name (str): 标签名称
            color (str): 标签颜色（HEX格式）
            
        Returns:
            bool: 是否添加成功
        """
        if not self.current_dataset:
            self.logger.error("没有打开的数据集")
            return False
        
        # 检查标签是否已存在
        for label in self.current_dataset["labels"]:
            if label["name"] == name:
                self.logger.warning(f"标签已存在：{name}")
                return False
        
        # 创建新标签
        label = {
            "id": len(self.current_dataset["labels"]),
            "name": name,
            "color": color,
            "count": 0
        }
        
        # 添加到数据集
        self.current_dataset["labels"].append(label)
        self._save_metadata()
        
        self.logger.info(f"已添加标签：{name}")
        return True
    
    def update_label_color(self, name, color):
        """
        更新标签颜色
        
        Args:
            name (str): 标签名称
            color (str): 新的颜色（HEX格式）
            
        Returns:
            bool: 是否更新成功
        """
        if not self.current_dataset:
            self.logger.error("没有打开的数据集")
            return False
        
        # 查找标签
        found = False
        for label in self.current_dataset["labels"]:
            if label["name"] == name:
                label["color"] = color
                found = True
                break
        
        if not found:
            self.logger.warning(f"找不到标签：{name}")
            return False
        
        # 保存更改
        self._save_metadata()
        self.logger.info(f"已更新标签颜色：{name} -> {color}")
        return True
    
    def rename_label(self, old_name, new_name):
        """
        重命名标签
        
        Args:
            old_name (str): 当前标签名称
            new_name (str): 新的标签名称
            
        Returns:
            bool: 是否重命名成功
        """
        if not self.current_dataset:
            self.logger.error("没有打开的数据集")
            return False
        
        # 检查新名称是否已存在
        for label in self.current_dataset["labels"]:
            if label["name"] == new_name:
                self.logger.warning(f"标签名称已存在：{new_name}")
                return False
        
        # 查找标签
        found = False
        for label in self.current_dataset["labels"]:
            if label["name"] == old_name:
                label["name"] = new_name
                found = True
                break
        
        if not found:
            self.logger.warning(f"找不到标签：{old_name}")
            return False
        
        # 保存更改
        self._save_metadata()
        self.logger.info(f"已重命名标签：{old_name} -> {new_name}")
        return True
    
    def remove_label(self, name):
        """
        删除标签
        
        Args:
            name (str): 标签名称
            
        Returns:
            bool: 是否删除成功
        """
        if not self.current_dataset:
            self.logger.error("没有打开的数据集")
            return False
        
        # 查找标签索引
        found_idx = -1
        for i, label in enumerate(self.current_dataset["labels"]):
            if label["name"] == name:
                found_idx = i
                break
        
        if found_idx == -1:
            self.logger.warning(f"找不到标签：{name}")
            return False
        
        # 从列表中移除
        removed_label = self.current_dataset["labels"].pop(found_idx)
        self.logger.info(f"已删除标签：{name}")
        
        # 重新排序剩余标签的ID
        for i, label in enumerate(self.current_dataset["labels"]):
            label["id"] = i
        
        # 保存更改
        self._save_metadata()
        return True
    
    def get_image_paths(self):
        """
        获取数据集中所有图像的路径
        
        Returns:
            list: 图像路径列表
        """
        if not self.current_dataset:
            self.logger.error("没有打开的数据集")
            return []
        
        dataset_path = Path(self.current_dataset["path"])
        images_dir = dataset_path / "images"
        
        image_paths = []
        for item in images_dir.iterdir():
            if item.is_file() and item.suffix.lower() in ['.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff']:
                image_paths.append(str(item))
        
        return image_paths
    
    def _save_metadata(self):
        """保存当前数据集的元数据"""
        if not self.current_dataset:
            return
        
        dataset_path = Path(self.current_dataset["path"])
        metadata_file = dataset_path / "dataset_meta.json"
        
        # 创建元数据的副本，移除path字段（这是运行时信息）
        metadata = dict(self.current_dataset)
        if "path" in metadata:
            metadata.pop("path")
        
        with open(metadata_file, "w", encoding="utf-8") as f:
            json.dump(metadata, f, ensure_ascii=False, indent=2) 